// Copyright (c) Aptos Foundation
// Licensed pursuant to the Innovation-Enabling Source Code License, available at https://github.com/aptos-labs/aptos-core/blob/main/LICENSE

#![allow(clippy::needless_borrow)]
#![allow(clippy::ptr_arg)]
#![allow(clippy::let_and_return)]

//! PVSS scheme-independent testing
#[cfg(test)]
use aptos_crypto::SecretSharingConfig;
use aptos_crypto::{
    blstrs::{random_scalar, G1_PROJ_NUM_BYTES, G2_PROJ_NUM_BYTES},
    weighted_config::WeightedConfigArkworks,
};
#[cfg(test)]
use aptos_dkg::pvss::traits::AggregatableTranscript;
use aptos_dkg::pvss::{
    chunky, das,
    das::unweighted_protocol,
    insecure_field, signed, test_utils,
    test_utils::{
        get_threshold_configs_for_benchmarking, get_weighted_configs_for_benchmarking,
        reconstruct_dealt_secret_key_randomly, NoAux,
    },
    traits::transcript::{
        HasAggregatableSubtranscript, NonAggregatableTranscript, Transcript, WithMaxNumShares,
    },
    GenericWeighting, ThresholdConfigBlstrs,
};
use rand::{rngs::StdRng, thread_rng};
use rand_core::SeedableRng;

// TODO: Add a test for public parameters serialization roundtrip?

#[test]
fn test_pvss_all_unweighted() {
    let mut rng = thread_rng();

    //
    // Unweighted PVSS tests
    //
    // TODO: might be better to grab (t,n) pairs, then intialise them for each PVSS using `new()` ? and consider renaming that `new_from_threshold()` ?
    let tcs = test_utils::get_threshold_configs_for_testing();
    for tc in tcs {
        println!("\nTesting {tc} PVSS");

        let seed = random_scalar(&mut rng);

        // Das
        pvss_deal_verify_and_reconstruct::<das::Transcript>(&tc, seed.to_bytes_le());

        // Insecure testing-only field-element PVSS
        pvss_deal_verify_and_reconstruct::<insecure_field::Transcript>(&tc, seed.to_bytes_le());
    }

    // Restarting the loop here because now it'll grab **arkworks** `ThresholdConfig`s over BN254 instead
    let tcs = test_utils::get_threshold_configs_for_testing_smaller();
    for tc in tcs.iter().take(20) {
        // Reduce the number of tcs to make it a bit faster?
        println!("\nTesting {tc} PVSS");

        let seed = random_scalar(&mut rng);

        type ChunkyTranscript = chunky::Transcript<ark_bn254::Bn254>;

        // Chunky
        pvss_nonaggregate_deal_verify_and_reconstruct::<chunky::Transcript<ark_bn254::Bn254>>(
            &tc,
            seed.to_bytes_le(),
        );

        pvss_deal_verify_and_reconstruct_from_subtranscript::<
            <ChunkyTranscript as Transcript>::SecretSharingConfig,
            ChunkyTranscript,
        >(&tc, seed.to_bytes_le());
    }
}

#[test]
fn test_pvss_all_weighted() {
    let mut rng = thread_rng();

    //
    // PVSS weighted tests
    //
    let wcs = test_utils::get_weighted_configs_for_testing();

    for wc in wcs {
        println!("\nTesting {wc} PVSS");
        let seed = random_scalar(&mut rng);

        // Generically-weighted Das
        // WARNING: Insecure, due to encrypting different shares with the same randomness, do not use!
        pvss_deal_verify_and_reconstruct::<GenericWeighting<das::Transcript>>(
            &wc,
            seed.to_bytes_le(),
        );

        // Generically-weighted field-element PVSS
        // WARNING: Insecure, reveals the dealt secret and its shares.
        pvss_deal_verify_and_reconstruct::<GenericWeighting<insecure_field::Transcript>>(
            &wc,
            seed.to_bytes_le(),
        );

        // Provably-secure Das PVSS
        pvss_deal_verify_and_reconstruct::<das::WeightedTranscript>(&wc, seed.to_bytes_le());
    }

    // Restarting the loop here because now it'll grab **arkworks** weighted `ThresholdConfig`s over BN254 instead
    let wcs = test_utils::get_weighted_configs_for_testing();
    for wc in wcs {
        println!("\nTesting {wc} PVSS");
        let seed = random_scalar(&mut rng);

        pvss_nonaggregate_weighted_deal_verify_and_reconstruct::<
            ark_bn254::Bn254,
            signed::GenericSigning<chunky::WeightedTranscript<ark_bn254::Bn254>>,
        >(&wc, seed.to_bytes_le());

        pvss_nonaggregate_weighted_deal_verify_and_reconstruct::<
            ark_bn254::Bn254,
            chunky::WeightedTranscript<ark_bn254::Bn254>,
        >(&wc, seed.to_bytes_le());
    }
}

#[test]
fn test_pvss_transcript_size() {
    for sc in get_threshold_configs_for_benchmarking() {
        println!();
        let expected_size = expected_transcript_size::<das::Transcript>(&sc);
        let actual_size = actual_transcript_size::<das::Transcript>(&sc);

        print_transcript_size::<das::Transcript>("Expected", &sc, expected_size);
        print_transcript_size::<das::Transcript>("Actual", &sc, actual_size);
    }

    // Restarting the loop here because now it'll grab **arkworks** `ThresholdConfig`s with BN254
    // uses default chunk sizes, so probably want to modify this at some point to allow a wider range
    // Ideally should iterate over a vec of (t, n), not the actual threshold configs... but won't be a bottleneck
    for sc in get_threshold_configs_for_benchmarking().iter().take(1) {
        // Only trying 1 for now to keep tests fast (also the second one has the same n, which means it would yield the same size...)
        println!();
        let actual_size = actual_transcript_size::<chunky::Transcript<ark_bn254::Bn254>>(&sc);
        print_transcript_size::<chunky::Transcript<ark_bn254::Bn254>>(
            "Actual for BN254",
            &sc,
            actual_size,
        );
    }

    // Restarting so it grabs BLS12-381 instead of BN254... TODO: could get rid of this with some work
    for sc in get_threshold_configs_for_benchmarking().iter().take(1) {
        // Only trying 1 for now to keep tests fast (also the second one has the same n, which means it would yield the same size...)

        println!();
        let actual_size =
            actual_transcript_size::<chunky::Transcript<ark_bls12_381::Bls12_381>>(&sc);
        print_transcript_size::<chunky::Transcript<ark_bls12_381::Bls12_381>>(
            "Actual for BLS12_381",
            &sc,
            actual_size,
        );
    }

    for wc in get_weighted_configs_for_benchmarking() {
        let actual_size = actual_transcript_size::<das::Transcript>(wc.get_threshold_config());
        print_transcript_size::<das::Transcript>("Actual", wc.get_threshold_config(), actual_size);

        let actual_size = actual_transcript_size::<das::WeightedTranscript>(&wc);
        print_transcript_size::<das::WeightedTranscript>("Actual", &wc, actual_size);
    }
}

#[cfg(test)]
fn print_transcript_size<T: Transcript>(size_type: &str, sc: &T::SecretSharingConfig, size: usize) {
    let name = T::scheme_name();
    println!("{size_type:8} transcript size for {sc} {name}: {size} bytes");
}

//
// Helper functions
//

/// Basic viability test for a PVSS transcript (weighted or unweighted):
///  1. Deals a secret, creating a transcript
///  2. Verifies the transcript.
///  3. Ensures the a sufficiently-large random subset of the players can recover the dealt secret
#[cfg(test)]
fn pvss_deal_verify_and_reconstruct<T: AggregatableTranscript>(
    sc: &T::SecretSharingConfig,
    seed_bytes: [u8; 32],
) {
    // println!();
    // println!("Seed: {}", hex::encode(seed_bytes.as_slice()));
    let mut rng = StdRng::from_seed(seed_bytes);

    let d = test_utils::setup_dealing::<T, StdRng>(sc, &mut rng);

    // Test dealing
    let trx = T::deal(
        &sc,
        &d.pp,
        &d.ssks[0],
        &d.spks[0],
        &d.eks,
        &d.s,
        &NoAux,
        &sc.get_player(0),
        &mut rng,
    );
    trx.verify(&sc, &d.pp, &[d.spks[0].clone()], &d.eks, &[NoAux])
        .expect("PVSS transcript failed verification");

    // Test transcript (de)serialization
    let trx_deserialized = T::try_from(trx.to_bytes().as_slice())
        .expect("serialized transcript should deserialize correctly");

    assert_eq!(trx, trx_deserialized);
    if d.dsk != reconstruct_dealt_secret_key_randomly::<StdRng, T>(sc, &mut rng, &d.dks, trx, &d.pp)
    {
        panic!("Reconstructed SK did not match");
    }
}

#[cfg(test)]
fn pvss_nonaggregate_deal_verify_and_reconstruct<T: NonAggregatableTranscript>(
    sc: &T::SecretSharingConfig,
    seed_bytes: [u8; 32],
) {
    // println!();
    // println!("Seed: {}", hex::encode(seed_bytes.as_slice()));
    let mut rng = StdRng::from_seed(seed_bytes);

    let d = test_utils::setup_dealing::<T, StdRng>(sc, &mut rng);

    // Test dealing
    let trx = T::deal(
        &sc,
        &d.pp,
        &d.ssks[0],
        &d.spks[0],
        &d.eks,
        &d.s,
        &NoAux,
        &sc.get_player(0),
        &mut rng,
    );
    trx.verify(&sc, &d.pp, &[d.spks[0].clone()], &d.eks, &NoAux)
        .expect("PVSS transcript failed verification");

    // Test transcript (de)serialization
    let trx_deserialized = T::try_from(trx.to_bytes().as_slice())
        .expect("serialized transcript should deserialize correctly");

    assert_eq!(trx, trx_deserialized);
    if d.dsk != reconstruct_dealt_secret_key_randomly::<StdRng, T>(sc, &mut rng, &d.dks, trx, &d.pp)
    {
        panic!("Reconstructed SK did not match");
    }
}

use ark_ec::pairing::Pairing;
// TODO: merge this stuff
#[cfg(test)]
fn pvss_nonaggregate_weighted_deal_verify_and_reconstruct<
    E: Pairing,
    T: NonAggregatableTranscript<SecretSharingConfig = WeightedConfigArkworks<E::ScalarField>>,
>(
    sc: &WeightedConfigArkworks<E::ScalarField>,
    seed_bytes: [u8; 32],
) {
    // println!();
    // println!("Seed: {}", hex::encode(seed_bytes.as_slice()));
    let mut rng = StdRng::from_seed(seed_bytes);

    let d = test_utils::setup_dealing_weighted::<E::ScalarField, T, StdRng>(sc, &mut rng);

    // Test dealing
    let trx = T::deal(
        &sc,
        &d.pp,
        &d.ssks[0],
        &d.spks[0],
        &d.eks,
        &d.s,
        &NoAux,
        &sc.get_player(0),
        &mut rng,
    );
    trx.verify(&sc, &d.pp, &[d.spks[0].clone()], &d.eks, &NoAux)
        .expect("PVSS transcript failed verification");

    // Test transcript (de)serialization
    let trx_deserialized = T::try_from(trx.to_bytes().as_slice())
        .expect("serialized transcript should deserialize correctly");

    assert_eq!(trx, trx_deserialized);
    if d.dsk != reconstruct_dealt_secret_key_randomly::<StdRng, T>(sc, &mut rng, &d.dks, trx, &d.pp)
    {
        panic!("Reconstructed SK did not match");
    }
}

#[cfg(test)]
fn pvss_deal_verify_and_reconstruct_from_subtranscript<
    C: SecretSharingConfig,
    T: Transcript<SecretSharingConfig = C> + HasAggregatableSubtranscript<C>,
>(
    sc: &T::SecretSharingConfig,
    seed_bytes: [u8; 32],
) {
    // println!();
    // println!("Seed: {}", hex::encode(seed_bytes.as_slice()));

    use aptos_dkg::pvss::test_utils::reconstruct_dealt_secret_key_randomly_subtranscript;
    let mut rng = StdRng::from_seed(seed_bytes);

    let d = test_utils::setup_dealing::<T, StdRng>(sc, &mut rng);

    // Test dealing
    let trx = T::deal(
        &sc,
        &d.pp,
        &d.ssks[0],
        &d.spks[0],
        &d.eks,
        &d.s,
        &NoAux,
        &sc.get_player(0),
        &mut rng,
    );

    let trx = trx.get_subtranscript();

    if d.dsk
        != reconstruct_dealt_secret_key_randomly_subtranscript::<StdRng, T::SubTranscript>(
            sc, &mut rng, &d.dks, trx, &d.pp,
        )
    {
        panic!("Reconstructed SK did not match");
    }
}

#[cfg(test)]
fn actual_transcript_size<T: Transcript>(sc: &T::SecretSharingConfig) -> usize {
    let mut rng = thread_rng();

    let trx = T::generate(
        &sc,
        &T::PublicParameters::with_max_num_shares_for_generate(sc.get_total_num_shares()),
        &mut rng,
    );
    let actual_size = trx.to_bytes().len();

    actual_size
}

#[cfg(test)]
fn expected_transcript_size<T: Transcript<SecretSharingConfig = ThresholdConfigBlstrs>>(
    sc: &ThresholdConfigBlstrs,
) -> usize {
    if T::scheme_name() == unweighted_protocol::DAS_SK_IN_G1 {
        G2_PROJ_NUM_BYTES
            + (sc.get_total_num_players() + 1) * (G2_PROJ_NUM_BYTES + G1_PROJ_NUM_BYTES)
    } else {
        panic!("Did not implement support for '{}' yet", T::scheme_name())
    }
}
